/*
 * <<
 *  ==
 *  Copyright (C) 2016 - 2019 EDP
 *  ==
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *        http://www.apache.org/licenses/LICENSE-2.0
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 *  >>
 *
 */

package com.sdkj.dspservice.modules.report.utils;

import com.sdkj.dspservice.common.constant.Consts;
import com.sdkj.dspservice.common.enums.DataTypeEnum;
import com.sdkj.dspservice.common.utils.DateUtils;
import com.sdkj.dspservice.common.utils.MD5Utils;
import com.sdkj.dspservice.common.utils.StringUtils;
import com.sdkj.dspservice.modules.report.entity.tem.PaginateWithQueryColumns;
import com.sdkj.dspservice.modules.report.entity.tem.QueryColumn;
import com.sdkj.dspservice.modules.report.entity.tem.TableInfo;
import com.sdkj.dspservice.modules.report.utils.entity.JdbcDataSource;
import com.sdkj.dspservice.modules.report.utils.entity.JdbcSourceInfo;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.Alias;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cache.annotation.CachePut;
import org.springframework.context.annotation.Scope;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import javax.sql.DataSource;
import java.rmi.ServerException;
import java.sql.*;
import java.text.MessageFormat;
import java.util.Date;
import java.util.*;
import java.util.regex.Matcher;

import static com.sdkj.dspservice.common.constant.Consts.*;
import static com.sdkj.dspservice.common.enums.DataTypeEnum.*;

/**
 * sql工具类
 *
 * @author tm
 */
@Slf4j
@Component
@Scope("prototype")
@Data
public class SqlUtils {

    private static final String TABLE = "TABLE";
    private static final String VIEW = "VIEW";
    private static final String[] TABLE_TYPES = new String[]{TABLE, VIEW};
    private static final String TABLE_NAME = "TABLE_NAME";
    private static final String TABLE_TYPE = "TABLE_TYPE";

    private int resultLimit;
    private boolean isQueryLogEnable;
    private JdbcSourceInfo jdbcSourceInfo;
    private SourceUtils sourceUtils;
    /**
     * 数据源类型
     */
    private DataTypeEnum dataTypeEnum;
    /**
     * 通过注入的方式实例化JdbcDataSource
     */
    @Autowired
    private JdbcDataSource jdbcDataSource;

    /**
     * 过滤sql中的注释
     *
     * @param sql
     * @return
     */
    public static String filterAnnotate(String sql) {
        sql = PATTERN_SQL_ANNOTATE.matcher(sql).replaceAll("$1");
        sql = sql.replaceAll(NEW_LINE_CHAR, SPACE).replaceAll("(;+\\s*)+", SEMICOLON);
        return sql;
    }

    /**
     * 检查敏感操作
     *
     * @param sql
     * @throws ServerException
     */
    public static void checkSensitiveSql(String sql) {
        Matcher matcher = PATTERN_SENSITIVE_SQL.matcher(sql.toLowerCase());
        if (matcher.find()) {
            String group = matcher.group();
            log.warn("Sensitive SQL operations are not allowed: {}", group.toUpperCase());
//            throw new ServerException("Sensitive SQL operations are not allowed: " + group.toUpperCase());
        }
    }

    public static boolean isSelect(String src) {
        if (com.alibaba.druid.util.StringUtils.isEmpty(src)) {
            return false;
        }
        try {
            Statement parse = CCJSqlParserUtil.parse(src);
            return parse instanceof Select;
        } catch (JSQLParserException e) {
            return false;
        }
    }

    public static Set<String> getQueryFromsAndJoins(String sql) {
        Set<String> columnPrefixes = new HashSet<>();
        try {
            Statement parse = CCJSqlParserUtil.parse(sql);
            Select select = (Select) parse;
            SelectBody selectBody = select.getSelectBody();
            if (selectBody instanceof PlainSelect) {
                PlainSelect plainSelect = (PlainSelect) selectBody;
                columnPrefixExtractor(columnPrefixes, plainSelect);
            }

            if (selectBody instanceof SetOperationList) {
                SetOperationList setOperationList = (SetOperationList) selectBody;
                List<SelectBody> selects = setOperationList.getSelects();
                for (SelectBody optSelectBody : selects) {
                    PlainSelect plainSelect = (PlainSelect) optSelectBody;
                    columnPrefixExtractor(columnPrefixes, plainSelect);
                }
            }

            if (selectBody instanceof WithItem) {
                WithItem withItem = (WithItem) selectBody;
                PlainSelect plainSelect = (PlainSelect) withItem.getSelectBody();
                columnPrefixExtractor(columnPrefixes, plainSelect);
            }
        } catch (JSQLParserException e) {
            log.debug(e.getMessage(), e);
        }
        return columnPrefixes;
    }

    private static void columnPrefixExtractor(Set<String> columnPrefixes, PlainSelect plainSelect) {
        getFromItemName(columnPrefixes, plainSelect.getFromItem());
        List<Join> joins = plainSelect.getJoins();
        if (!CollectionUtils.isEmpty(joins)) {
            joins.forEach(join -> getFromItemName(columnPrefixes, join.getRightItem()));
        }
    }

    private static void getFromItemName(Set<String> columnPrefixs, FromItem fromItem) {
        if (fromItem == null) {
            return;
        }
        Alias alias = fromItem.getAlias();
        if (alias != null) {
            if (alias.isUseAs()) {
                columnPrefixs.add(alias.getName().trim() + DOT);
            } else {
                columnPrefixs.add(alias.toString().trim() + DOT);
            }
        } else {
            fromItem.accept(getFromItemTableName(columnPrefixs));
        }
    }

    private static FromItemVisitor getFromItemTableName(Set<String> set) {
        return new FromItemVisitor() {
            @Override
            public void visit(Table tableName) {
                set.add(tableName.getName() + DOT);
            }

            @Override
            public void visit(SubSelect subSelect) {
            }

            @Override
            public void visit(SubJoin subjoin) {
            }

            @Override
            public void visit(LateralSubSelect lateralSubSelect) {
            }

            @Override
            public void visit(ValuesList valuesList) {
            }

            @Override
            public void visit(TableFunction tableFunction) {
            }

            @Override
            public void visit(ParenthesisFromItem aThis) {
            }
        };
    }

    /**
     * 获取查询行数的SQL
     *
     * @param sql
     * @return
     */
    public static String getCountSql(String sql) {
        String countSql = String.format(Consts.QUERY_COUNT_SQL, sql);
        try {
            Select select = (Select) CCJSqlParserUtil.parse(sql);
            PlainSelect plainSelect = (PlainSelect) select.getSelectBody();
            plainSelect.setOrderByElements(null);
            countSql = String.format(QUERY_COUNT_SQL, select.toString());
        } catch (JSQLParserException e) {
            log.debug(e.getMessage(), e);
        }
        return SqlParseUtils.rebuildSqlWithFragment(countSql);
    }

    public static String getColumnLabel(Set<String> columnPrefixs, String columnLable) {
        if (!CollectionUtils.isEmpty(columnPrefixs)) {
            for (String prefix : columnPrefixs) {
                if (columnLable.startsWith(prefix)) {
                    return columnLable.replaceFirst(prefix, EMPTY);
                }
                if (columnLable.startsWith(prefix.toLowerCase())) {
                    return columnLable.replaceFirst(prefix.toLowerCase(), EMPTY);
                }
                if (columnLable.startsWith(prefix.toUpperCase())) {
                    return columnLable.replaceFirst(prefix.toUpperCase(), EMPTY);
                }
            }
        }
        return columnLable;
    }

    /**
     * 初始化（这里不用构造器的原因是，JdbcDataSource需要单例，用的是注入方式，作为JdbcDataSource的父类也必须注入方式，否则jdbcDataSource为null）
     *
     * @param jdbcSourceInfo
     */
    public SqlUtils init(JdbcSourceInfo jdbcSourceInfo) {
        SqlUtils sqlUtils = new SqlUtils();
        sqlUtils.setJdbcSourceInfo(jdbcSourceInfo);
        sqlUtils.setSourceUtils(new SourceUtils(this.jdbcDataSource));
        sqlUtils.setJdbcDataSource(this.jdbcDataSource);
        return sqlUtils;
    }

    /**
     * 获取当前连接数据库
     *
     * @return
     * @throws SourceException
     */
    public String getDatabaseName() throws SourceException {
        Connection connection = null;
        try {
            connection = sourceUtils.getConnection(this.jdbcSourceInfo);
            if (null == connection) {
                return null;
            }

            if (dataTypeEnum == ORACLE) {
                return this.jdbcSourceInfo.getUsername();
            }

            if (dataTypeEnum == ELASTICSEARCH) {
                if (StringUtils.isEmpty(this.jdbcSourceInfo.getUsername())) {
                    return dataTypeEnum.getFeature();
                } else {
                    return this.jdbcSourceInfo.getUsername();
                }
            }
            // 获取当前数据库
            return connection.getCatalog();
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            return null;
        } finally {
            SourceUtils.releaseConnection(connection);
        }
    }

    /**
     * 获取当前连接数据库列表
     *
     * @return
     * @throws SourceException
     */
    public List<String> getDatabases() throws SourceException {
        List<String> dbList = new ArrayList<>();
        Connection connection = null;
        try {
            connection = sourceUtils.getConnection(this.jdbcSourceInfo);
            if (null == connection) {
                return dbList;
            }

            if (dataTypeEnum == ORACLE) {
                dbList.add(this.jdbcSourceInfo.getUsername());
                return dbList;
            }

            if (dataTypeEnum == ELASTICSEARCH) {
                if (StringUtils.isEmpty(this.jdbcSourceInfo.getUsername())) {
                    dbList.add(dataTypeEnum.getFeature());
                } else {
                    dbList.add(this.jdbcSourceInfo.getUsername());
                }
                return dbList;
            }
            // 获取数据库列表
            DatabaseMetaData metaData = connection.getMetaData();
            ResultSet rs = metaData.getCatalogs();
            while (rs.next()) {
                dbList.add(rs.getString(1));
            }
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            return dbList;
        } finally {
            SourceUtils.releaseConnection(connection);
        }
        return dbList;
    }

    /**
     * 获取数据库的数据表名称列表
     *
     * @param dbName 数据库名称
     * @return
     * @throws SourceException
     */
    public List<String> getTableNameList(String dbName) throws SourceException {

        List<String> tableList = null;
        Connection connection = null;
        ResultSet tables = null;

        try {
            connection = sourceUtils.getConnection(this.jdbcSourceInfo);
            if (null == connection) {
                return null;
            }
            DatabaseMetaData metaData = connection.getMetaData();
            String schema = null;
            try {
                schema = metaData.getConnection().getSchema();
            } catch (Throwable t) {
                // ignore
            }

            tables = metaData.getTables(dbName, getDBSchemaPattern(schema), "%", TABLE_TYPES);
            if (null == tables) {
                return null;
            }

            tableList = new ArrayList<>();
            while (tables.next()) {
                String name = tables.getString(TABLE_NAME);
                if (!com.alibaba.druid.util.StringUtils.isEmpty(name)) {
//                    String type = TABLE;
//                    try {
//                        type = tables.getString(TABLE_TYPE);
//                    } catch (Exception e) {
//                        // ignore
//                    }
//                    tableList.add(new QueryColumn(name, type));
                    tableList.add(name);
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
            log.error(e.toString(), e);
            return tableList;
        } finally {
            SourceUtils.closeResult(tables);
            SourceUtils.releaseConnection(connection);
        }
        return tableList;
    }

    /**
     * 获取指定表列信息
     *
     * @param tableName
     * @return
     * @throws SourceException
     */
    public TableInfo getTableInfo(String dbName, String tableName) throws SourceException {
        TableInfo tableInfo = null;
        Connection connection = null;
        try {
            connection = sourceUtils.getConnection(this.jdbcSourceInfo);
            if (null != connection) {
                DatabaseMetaData metaData = connection.getMetaData();
                List<String> primaryKeys = getPrimaryKeys(dbName, tableName, metaData);
                List<QueryColumn> columns = getColumns(dbName, tableName, metaData);
                tableInfo = new TableInfo(tableName, primaryKeys, columns);
            }
        } catch (SQLException e) {
            log.error(e.toString(), e);
            throw new SourceException(e.getMessage() + ", jdbcUrl=" + this.jdbcSourceInfo.getJdbcUrl());
        } finally {
            SourceUtils.releaseConnection(connection);
        }
        return tableInfo;
    }

    /**
     * 获取数据表主键
     *
     * @param tableName
     * @param metaData
     * @return
     */
    private List<String> getPrimaryKeys(String dbName, String tableName, DatabaseMetaData metaData) {
        ResultSet rs = null;
        List<String> primaryKeys = new ArrayList<>();
        try {
            rs = metaData.getPrimaryKeys(dbName, null, tableName);
            if (rs == null) {
                return primaryKeys;
            }
            while (rs.next()) {
                primaryKeys.add(rs.getString("COLUMN_NAME"));
            }
        } catch (Exception e) {
            log.error(e.getMessage(), e);
        } finally {
            SourceUtils.closeResult(rs);
        }
        return primaryKeys;
    }

    /**
     * 获取数据表列
     *
     * @param tableName
     * @param metaData
     * @return
     */
    private List<QueryColumn> getColumns(String dbName, String tableName, DatabaseMetaData metaData) {
        ResultSet rs = null;
        List<QueryColumn> columnList = new ArrayList<>();
        try {
            if (this.dataTypeEnum == ORACLE) {
                dbName = null;
            }
            rs = metaData.getColumns(dbName, null, tableName, "%");
            if (rs == null) {
                return columnList;
            }
            while (rs.next()) {
                columnList.add(new QueryColumn(rs.getString("COLUMN_NAME"), rs.getString("TYPE_NAME"), rs.getString("REMARKS")));
            }
        } catch (Exception e) {
            log.error(e.getMessage(), e);
        } finally {
            SourceUtils.closeResult(rs);
        }
        return columnList;
    }

    private String getDBSchemaPattern(String schema) {
        if (dataTypeEnum == null) {
            return null;
        }
        String schemaPattern = null;
        switch (dataTypeEnum) {
            case ORACLE:
                schemaPattern = this.jdbcSourceInfo.getUsername();
                if (null != schemaPattern) {
                    schemaPattern = schemaPattern.toUpperCase();
                }
                break;
            case SQLSERVER:
                schemaPattern = "dbo";
                break;
            case CLICKHOUSE:
            case PRESTO:
                if (!com.alibaba.druid.util.StringUtils.isEmpty(schema)) {
                    schemaPattern = schema;
                }
                break;
            default:
                break;
        }
        return schemaPattern;
    }

    public JdbcTemplate jdbcTemplate() throws SourceException {
        Connection connection = null;
        try {
            connection = sourceUtils.getConnection(jdbcSourceInfo);
        } finally {
            SourceUtils.releaseConnection(connection);
        }
        DataSource dataSource = sourceUtils.getJdbcDataSource().getDataSource(jdbcSourceInfo);
        JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource);
        jdbcTemplate.setFetchSize(500);
        return jdbcTemplate;
    }

    public void execute(String sql) {
        sql = filterAnnotate(sql);
        checkSensitiveSql(sql);
        if (isQueryLogEnable) {
            String md5 = MD5Utils.getMD5(sql, true, 16);
//            sqlLogger.info("{} execute for sql:{}", md5, formatSql(sql));
        }
        try {
            jdbcTemplate().execute(sql);
        } catch (Exception e) {
            log.error(e.getMessage(), e);
        }
    }

    /**
     * 执行SQL
     *
     * @param sql
     * @param excludeColumns
     * @return
     */
    public PaginateWithQueryColumns executeSql(String sql, int pageNo, int pageSize) {
        return executeSql(sql, pageNo, pageSize, 0, -1, null);
    }

    /**
     * 执行SQL
     *
     * @param sql
     * @param excludeColumns
     * @return
     */
    @CachePut(value = "query", keyGenerator = "keyGenerator")
    public PaginateWithQueryColumns executeSql(String sql, int pageNo, int pageSize, int totalCount, int limit, Set<String> excludeColumns) {
        PaginateWithQueryColumns paginateWithQueryColumns = new PaginateWithQueryColumns();
        sql = filterAnnotate(sql);
        checkSensitiveSql(sql);

        long before = System.currentTimeMillis();

        JdbcTemplate jdbcTemplate = jdbcTemplate();
        jdbcTemplate.setMaxRows(resultLimit);
        if (pageNo < 1 && pageSize < 1) {
            if (limit > 0) {
                jdbcTemplate.setMaxRows(Math.min(limit, resultLimit));
            }
            // special for mysql
            if (getDataTypeEnum() == MYSQL) {
                jdbcTemplate.setFetchSize(Integer.MIN_VALUE);
            }

            getResultForPaginate(sql, paginateWithQueryColumns, jdbcTemplate, excludeColumns, -1);
            paginateWithQueryColumns.setPageNo(1);
            int size = paginateWithQueryColumns.getResultList().size();
            paginateWithQueryColumns.setPageSize(size);
            paginateWithQueryColumns.setTotalCount(size);

        } else {
            paginateWithQueryColumns.setPageNo(pageNo);
            paginateWithQueryColumns.setPageSize(pageSize);
            // 计算开始行数，起始为0，含头不含尾
            int startRow = (pageNo - 1) * pageSize;
            // 当为第一页或者总数为0时，则去获取总数
            if (pageNo == 1 || totalCount == 0) {
                // 获取总数
                Object o = jdbcTemplate.queryForList(getCountSql(sql), Object.class).get(0);
                totalCount = Integer.parseInt(String.valueOf(o));
            }

            if (limit > 0) {
                totalCount = Math.min(Math.min(limit, resultLimit), totalCount);
                if (limit < pageNo * pageSize) {
                    jdbcTemplate.setMaxRows(limit - startRow);
                } else {
                    jdbcTemplate.setMaxRows(Math.min(limit, pageSize));
                }
            } else {
                // 计算结束行数，起始为0，含头不含尾
                jdbcTemplate.setMaxRows(pageNo * pageSize);
            }

            paginateWithQueryColumns.setTotalCount(totalCount);
            // 获取结果
            if (this.dataTypeEnum == MYSQL) {
                sql = sql + " LIMIT " + startRow + ", " + pageSize;
                getResultForPaginate(sql, paginateWithQueryColumns, jdbcTemplate, excludeColumns, -1);
            } else {
                getResultForPaginate(sql, paginateWithQueryColumns, jdbcTemplate, excludeColumns, startRow);
            }
        }
        // 打印日志
//        if (isQueryLogEnable) {
        String md5 = MD5Utils.getMD5(sql + pageNo + pageSize + limit, true, 16);
        System.out.println("************************************************************************");
        System.out.println(MessageFormat.format("{0} query for({1} ms) total count: {2}, page size: {3}, sql:{4}",
                md5, System.currentTimeMillis() - before,
                paginateWithQueryColumns.getTotalCount(),
                paginateWithQueryColumns.getPageSize(), sql));
        System.out.println("************************************************************************");
//        }

        return paginateWithQueryColumns;
    }

    /**
     * 获取SQL查询结果，即执行jdbcTemplate.query
     *
     * @param sql
     * @param paginateWithQueryColumns
     * @param jdbcTemplate
     * @param excludeColumns
     * @param startRow
     */
    private void getResultForPaginate(String sql, PaginateWithQueryColumns paginateWithQueryColumns, JdbcTemplate jdbcTemplate, Set<String> excludeColumns, int startRow) {
        Set<String> queryFromsAndJoins = getQueryFromsAndJoins(sql);
        jdbcTemplate.query(sql, rs -> {
            if (null == rs) {
                return paginateWithQueryColumns;
            }

            ResultSetMetaData metaData = rs.getMetaData();
            List<QueryColumn> queryColumns = new ArrayList<>();
            for (int i = 1; i <= metaData.getColumnCount(); i++) {
                String key = getColumnLabel(queryFromsAndJoins, metaData.getColumnLabel(i));
                if (!CollectionUtils.isEmpty(excludeColumns) && excludeColumns.contains(key)) {
                    continue;
                }
                queryColumns.add(new QueryColumn(key, metaData.getColumnTypeName(i)));
            }
            paginateWithQueryColumns.setColumns(queryColumns);

            List<Map<String, Object>> resultList = new ArrayList<>();

            try {
                if (startRow > 0) {
                    rs.absolute(startRow);
                }
                while (rs.next()) {
                    resultList.add(getResultObjectMap(excludeColumns, rs, metaData, queryFromsAndJoins));
                }
            } catch (Throwable e) {
                int currentRow = 0;
                while (rs.next()) {
                    if (currentRow >= startRow) {
                        resultList.add(getResultObjectMap(excludeColumns, rs, metaData, queryFromsAndJoins));
                    }
                    currentRow++;
                }
            }

            paginateWithQueryColumns.setResultList(resultList);

            return paginateWithQueryColumns;
        });
    }

    private Map<String, Object> getResultObjectMap(Set<String> excludeColumns, ResultSet rs, ResultSetMetaData metaData, Set<String> queryFromsAndJoins) throws SQLException {
        Map<String, Object> map = new LinkedHashMap<>();

        for (int i = 1; i <= metaData.getColumnCount(); i++) {
            String key = metaData.getColumnLabel(i);
            String label = getColumnLabel(queryFromsAndJoins, key);

            if (!CollectionUtils.isEmpty(excludeColumns) && excludeColumns.contains(label)) {
                continue;
            }
            Object value = rs.getObject(key);
            // 如果是日期格式，转为字符串
            if(value instanceof Timestamp){
                value = DateUtils.formatDate(new Date(((Timestamp)rs.getObject(key)).getTime()), "yyyy-MM-dd HH:mm:ss");
            }
            map.put(label, value instanceof byte[] ? new String((byte[]) value) : value);
        }
        return map;
    }

}

